'''
将处理好的BTCV和CHAOS复制到另外一个文件夹，并划分好训练集和测试集
'''
import os

from sklearn.model_selection import train_test_split
import shutil

def intTostr(n):
    '''
    将数字转化为三位小数，在前面补0
    '''
    if n<10:
        return '00'+str(n)
    if 9<n<100:
        return '0' + str(n)
    if 99<n<1000:
        return  str(n)
    return str(n)
def copy(dataroot="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA",
         source="BTCV",target="CHAOS",
         save_dir="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/useData"):
    '''
    将source复制一份到save_dir
    将target分为target和test两个文件夹（80%和20%）
    '''

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    if os.path.exists(os.path.join(save_dir,"source")):#如果文件夹已存在，则删除
        shutil.rmtree(os.path.join(save_dir,"source"))
    if os.path.exists(os.path.join(save_dir,"target")):#如果文件夹已存在，则删除
        shutil.rmtree(os.path.join(save_dir,"target"))
    if os.path.exists(os.path.join(save_dir,"test")):#如果文件夹已存在，则删除
        shutil.rmtree(os.path.join(save_dir,"test"))

    if not os.path.exists(os.path.join(save_dir,"source")):
        os.makedirs(os.path.join(save_dir,"source"))
    if not os.path.exists(os.path.join(save_dir,"target")):
        os.makedirs(os.path.join(save_dir,"target"))
    if not os.path.exists(os.path.join(save_dir,"test")):
        os.makedirs(os.path.join(save_dir,"test"))
    # --------复制source-----#
    source_number = 0
    for subject in os.listdir(os.path.join(dataroot,source)):
        subject_dir=os.path.join(dataroot,source,subject)
        for f in os.listdir(subject_dir):
            source_number+=1
            p_source=os.path.join(dataroot,source,subject,f)

            p_target = os.path.join(save_dir, "source", subject+'_'+f)

            shutil.copyfile(p_source,p_target)
    #-------划分target的训练和测试-------#
    target_number=0
    target_subject_list=os.listdir(os.path.join(dataroot,target))
    train_list,test_list=train_test_split(target_subject_list,test_size=0.2,random_state=100)
    for subject in train_list:
        subject_dir = os.path.join(dataroot, target, subject)
        for f in os.listdir(subject_dir):
            target_number += 1
            p_source = os.path.join(dataroot, target, subject, f)
            p_target = os.path.join(save_dir, "target", subject+'_'+f)
            shutil.copyfile(p_source, p_target)

    test_number = 0
    for subject in test_list:
        subject_dir = os.path.join(dataroot, target, subject)
        for f in os.listdir(subject_dir):
            test_number += 1
            p_source = os.path.join(dataroot, target, subject, f)
            p_target = os.path.join(save_dir, "test", subject+'_'+f)
            shutil.copyfile(p_source, p_target)
if __name__ == '__main__':
    copy(dataroot="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA",
         source="BTCV4label", target="CHAOS4label",
         save_dir="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/useData4label")